{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import wandb\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "\n",
    "def mean_and_std(df):\n",
    "    agg = np.stack(df.to_numpy(), axis=0)\n",
    "    return np.mean(agg, axis=0), np.std(agg, axis=0)\n",
    "\n",
    "download_root = \".\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sweep_regression_df_all(sweep_id, allow_crash=False):\n",
    "    api = wandb.Api()\n",
    "    sweep = api.sweep(\"ngruver/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in sweep.runs:        \n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "        \n",
    "        if \"finished\" in str(run):\n",
    "            summary = pd.Series(run.summary)\n",
    "        else:\n",
    "            history = run.history()\n",
    "            summary = pd.Series({k: history[k].to_numpy()[-1] for k,v in history.items()})\n",
    "        results.append(pd.concat([config,summary]))\n",
    "    return pd.concat(results,axis=1).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = get_sweep_regression_df_all(\"v96kirjy\",allow_crash=True)\n",
    "df2 = get_sweep_regression_df_all(\"pexiwka8\",allow_crash=True)\n",
    "df = pd.concat((df,df2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"model_type\"].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "reader_friendly_dict = {\n",
    "    \"NN\": \"NODE\",\n",
    "    \"MechanicsNN\": \"NODE + SO\",\n",
    "    \"HNN\": \"HNN\"\n",
    "}\n",
    "\n",
    "sns.set_style('whitegrid')\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "\n",
    "filtered = df[df['model_type'].isin(['HNN','NN','MechanicsNN'])].copy()\n",
    "filtered[\"model_type\"] = filtered[\"model_type\"].apply(lambda s: reader_friendly_dict[s])\n",
    "filtered[\"dataset\"]=filtered[\"system_type\"]+filtered[\"num_bodies\"].astype(str)\n",
    "filtered[\"Rollout Error\"] = filtered[\"test_gerr\"].astype(float)\n",
    "filtered[\"Energy Violation\"] = filtered[\"test_Herr\"].astype(float)\n",
    "filled_markers = ('o', 'v', '^', '<', '>', '8', 's', 'p', '*', 'h', 'H', 'D', 'd', 'P', 'X')[:len(filtered[\"dataset\"].unique())]\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'\n",
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4.5,3.5))\n",
    "sns.scatterplot(data=filtered,x='Rollout Error',y='Energy Violation',hue='model_type', ax=ax)#,style=\"dataset\",markers=filled_markers)\n",
    "ax.get_legend().remove()\n",
    "plt.yscale('log')\n",
    "plt.xscale('log')\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=3)\n",
    "fig.subplots_adjust(bottom=0.3)\n",
    "\n",
    "plt.savefig('energy_conservation_loglog.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import datasets, linear_model, metrics\n",
    "regr = linear_model.LinearRegression()\n",
    "regr.fit(np.log(filtered[\"Rollout Error\"][:,None]), np.log(filtered[\"Energy Violation\"]))\n",
    "y_pred = regr.predict(np.log(filtered[\"Rollout Error\"][:,None]))\n",
    "y_true = np.log(filtered[\"Energy Violation\"])\n",
    "residuals = y_true-y_pred\n",
    "filtered[\"residuals\"] = residuals/np.log(filtered[\"Energy Violation\"]).std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics.r2_score(y_true, y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(4, 3))\n",
    "order = sorted(filtered[\"dataset\"].unique())\n",
    "order = order[5:]+order[:5]\n",
    "\n",
    "plot =sns.barplot(y=\"residuals\",hue=\"model_type\",x=\"dataset\",data=filtered,order=order)\n",
    "plt.setp(plot.get_xticklabels(), rotation=30)\n",
    "plt.xlabel('')\n",
    "plt.savefig('energy_conservation_residuals.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = get_sweep_regression_df_all(\"kj4ke9i2\",allow_crash=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered = df#df[df['model_type'].str.fullmatch('|'.join(['HNN','NN','MechanicsNN','SecondOrderNN']))]\n",
    "filtered[\"dataset\"]=filtered[\"system_type\"].apply(lambda s: s.replace(\"Pendulum\", \" \")) +filtered[\"num_bodies\"].astype(str)\n",
    "\n",
    "filtered[\"SymReg strength\"] = 1/filtered[\"alpha\"].astype(float)\n",
    "order = sorted(filtered[\"dataset\"].unique())\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6.75,5.25))\n",
    "plot = sns.barplot(data=filtered, x=\"dataset\", y='test_gerr', hue=\"SymReg strength\", order=order, palette=\"rocket\",ax=ax)\n",
    "ax.get_legend().remove()\n",
    "ax.grid(False)\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('')\n",
    "plt.ylabel(\"Rollout Error\")\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "leg = fig.legend(handles, labels, loc='lower center', ncol=6, prop={'size': 12}, title=\"$\\\\alpha=$\")#, fontsize=45)\n",
    "fig.subplots_adjust(bottom=0.2, left=-.15)\n",
    "    \n",
    "plt.savefig('state_err_reg.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib.ticker import FuncFormatter\n",
    "\n",
    "filtered = df\n",
    "filtered[\"dataset\"]=filtered[\"system_type\"].apply(lambda s: s.replace(\"Pendulum\", \" \")) +filtered[\"num_bodies\"].astype(str)\n",
    "filtered[\"Symplectic Error\"] = np.log10(filtered[\"Train_symreg\"].astype(float))\n",
    "filtered[\"Rollout Error\"] = np.log10(filtered[\"test_gerr\"].astype(float))\n",
    "\n",
    "filtered[\"SymReg strength\"] = 1/filtered[\"alpha\"].astype(float)\n",
    "order = sorted(filtered[\"dataset\"].unique())\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'\n",
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "palette = sns.color_palette(\"Paired\", desat=0.8)[4:]\n",
    "g = sns.lmplot(data=filtered,x=\"Symplectic Error\",y='Rollout Error',hue=\"dataset\",hue_order=order, legend_out=True, height=4, aspect=1.3, palette=palette)\n",
    "\n",
    "ax = g.axes[0,0]\n",
    "ax.set_xticks(np.arange(-10,4,2))\n",
    "ax.set_yticks(np.arange(-4,1))\n",
    "formatter = lambda x, pos: f'{10. ** x:g}'\n",
    "ax.get_xaxis().set_major_formatter(FuncFormatter(formatter))\n",
    "ax.get_yaxis().set_major_formatter(FuncFormatter(formatter))\n",
    "\n",
    "ax.grid(False)\n",
    "ax.tick_params(axis='both', which='major', labelsize=14)\n",
    "ax.set_xlabel(\"Symplectic Error\", labelpad=10)\n",
    "\n",
    "legend = g.legend\n",
    "legend.set_title(\"\")\n",
    "\n",
    "plt.savefig('state_err_reg_value.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import wandb\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "\n",
    "def mean_and_std(df):\n",
    "    agg = np.stack(df.to_numpy(), axis=0)\n",
    "    return np.mean(agg, axis=0), np.std(agg, axis=0)\n",
    "\n",
    "download_root = \".\"\n",
    "import json\n",
    "def get_sweep_tabular(sweep_id, allow_crash=True):\n",
    "    api = wandb.Api()\n",
    "    sweep = api.sweep(\"ngruver/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in sweep.runs:\n",
    "#         print(run)\n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "        if \"finished\" in str(run):\n",
    "#             print(run.summary)\n",
    "            summary = pd.Series(run.summary)\n",
    "        else:\n",
    "            history = run.history()\n",
    "            summary = pd.Series({k: history[k].to_numpy()[-1] for k,v in history.items()})\n",
    "        for f in run.files():\n",
    "            if not f.name.endswith(summary['H_err_vec']['path'].split('/')[-1]):\n",
    "                continue\n",
    "            f.download(root=\".\", replace=True)\n",
    "            with open(f.name) as fd:\n",
    "                data = np.array(json.load(fd)['data'])\n",
    "                print(f.name)\n",
    "        config = pd.Series(run.config)\n",
    "        logherrs=data\n",
    "        ic = np.arange(logherrs.shape[0])[:,None]\n",
    "        ic =ic+ np.zeros_like(logherrs)\n",
    "        T = np.linspace(0,1,ic.shape[-1])[None,:]+np.zeros_like(logherrs)\n",
    "        df = pd.DataFrame({'logherr':logherrs.reshape(-1),'ics':ic.reshape(-1),'T':T.reshape(-1)})\n",
    "        c = config.to_frame()\n",
    "        for att in c.T.columns:\n",
    "            df[att]=config[att]\n",
    "        results.append(df)\n",
    "    return pd.concat(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_all = get_sweep_tabular(\"j3sjkwvo\",False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"dataset\"]=df[\"system_type\"]+df[\"num_bodies\"].astype(str)\n",
    "mean = df.groupby(['model_type','dataset','T']).mean()['logherr'].reset_index()\n",
    "std = df.groupby(['model_type','dataset','T']).std()['logherr'].reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean['std'] = std['logherr']\n",
    "mean['std'] = np.exp(mean['std'])\n",
    "mean['logherr']=np.exp(mean['logherr'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "\n",
    "fig1, f1_axes = plt.subplots(ncols=3, nrows=2, constrained_layout=True,figsize=(8,6),sharex=True,sharey=True)\n",
    "datasets = [f'ChainPendulum{i}' for i in (2,3,4)]+[f'SpringPendulum{i}' for i in (2,3,4)]\n",
    "for i,ds in enumerate(datasets):\n",
    "    dfs = mean[mean['dataset']==ds]\n",
    "    #print(dfs[dfs['T']==1])\n",
    "    dfhnn = dfs[dfs['model_type']=='HNN']\n",
    "    dfnn =  dfs[dfs['model_type']=='NN']\n",
    "    ax = f1_axes[i//3,i%3]\n",
    "    ax.plot(dfhnn['T'],dfhnn['logherr'],label=\"HNN\")\n",
    "    ax.fill_between(dfhnn['T'],dfhnn['logherr']/dfhnn['std'],dfhnn['logherr']*dfhnn['std'],alpha=.2)\n",
    "    ax.plot(dfnn['T'],dfnn['logherr'],label=\"NODE\")\n",
    "    ax.fill_between(dfnn['T'],dfnn['logherr']/dfnn['std'],dfnn['logherr']*dfnn['std'],alpha=.2)\n",
    "    ax.set_xscale('log')\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_ylim(bottom=5e-6, top=1e1)\n",
    "    ax.grid(True)\n",
    "    ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "    if i//3==0:\n",
    "        ax.set_title(f\"{i+2} link\")\n",
    "    #ax.title(ds.split('P')[0])\n",
    "fig1.text(1.01, 0.72, 'Chain', ha='center', va='center', rotation='vertical')\n",
    "fig1.text(1.01, 0.3, 'Spring', ha='center', va='center', rotation='vertical')\n",
    "fig1.text(-0.005, 0.5, 'Energy Error', ha='center', va='center', rotation='vertical')\n",
    "fig1.text(0.54, 0, 'Rollout Time T', ha='center', va='center')\n",
    "\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "fig1.savefig('energy_growth.pdf', bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
